#pragma once
#include "registration_interface.hpp"
#include "ndt_cpu/NormalDistributionsTransform.h"
#include <pcl/common/transforms.h>

namespace TESTICP
{

    class NDTCPURegistration : public RegistrationInterface
    { // 继承点云配准的基类
    public:
        NDTCPURegistration(const YAML::Node& node);

        void setInputTarget(const CloudPtr& target) override;
        void setInputSource(const CloudPtr& source) override;

        bool align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose = Eigen::Matrix4f::Identity()) override;

        inline Eigen::Matrix4f getFinalTransformation() { return ndt_cpu_.getFinalTransformation(); }
        inline float getFitnessScore() { return ndt_cpu_.getFitnessScore(); }

    private:
        bool SetRegistrationParam(float res, float step_size, float trans_eps, int max_iter);

    private:
        cpu::NormalDistributionsTransform<PointType, PointType> ndt_cpu_; // 实例化cpu_ndt 对象
    };

    NDTCPURegistration::NDTCPURegistration(const YAML::Node& node) {
        float res = node["res"].as<float>();
        float step_size = node["step_size"].as<float>();
        float trans_eps = node["trans_eps"].as<float>();
        int max_iter = node["max_iter"].as<int>();

        SetRegistrationParam(res, step_size, trans_eps, max_iter);
    }

    bool NDTCPURegistration::SetRegistrationParam(float res, float step_size, float trans_eps, int max_iter) {
        ndt_cpu_.setResolution(res);
        ndt_cpu_.setStepSize(step_size);
        ndt_cpu_.setTransformationEpsilon(trans_eps);
        ndt_cpu_.setMaximumIterations(max_iter);

        std::cout << "NDT params:" << std::endl
            << "res: " << res << ", "
            << "step_size: " << step_size << ", "
            << "trans_eps: " << trans_eps << ", "
            << "max_iter: " << max_iter
            << std::endl
            << std::endl;

        return true;
    }

    void NDTCPURegistration::setInputTarget(const CloudPtr& target) {
        ndt_cpu_.setInputTarget(target);
    }

    void NDTCPURegistration::setInputSource(const CloudPtr& source) {
        ndt_cpu_.setInputSource(source);
    }

    bool NDTCPURegistration::align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose) {
        ndt_cpu_.align(output_ptr, predict_pose);
        return true;
    }
}